[https://nvbugs/6546909][fix] Make the handle's lifetime match its graphs — lazy… - #17313
[https://nvbugs/6546909][fix] Make the handle's lifetime match its graphs — lazy…#17313trtllm-agent wants to merge 1 commit into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
WalkthroughThe GB300 skip waiver is removed for the Llama 3.1 8B FP8 integration test configuration using ChangesIntegration test waiver update
Estimated code review effort: 1 (Trivial) | ~2 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/compilation/backend.py (1)
124-138: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd docstrings for the graph-pool lifecycle methods.
get_graph_pool_handle()andretire_graph_pool_handle()are used across module boundaries. Replace the implementation comments with Google-style docstrings that define their lifecycle contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/compilation/backend.py` around lines 124 - 138, Replace the inline comments in get_graph_pool_handle() and retire_graph_pool_handle() with Google-style docstrings documenting lazy handle creation, replacement after retirement, and the requirement to retire only after captured graphs are reset. Preserve the existing implementation behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/compilation/backend.py`:
- Around line 124-138: Replace the inline comments in get_graph_pool_handle()
and retire_graph_pool_handle() with Google-style docstrings documenting lazy
handle creation, replacement after retirement, and the requirement to retire
only after captured graphs are reset. Preserve the existing implementation
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0b58bdab-da34-4805-bbb4-fc4b75c759e0
📒 Files selected for processing (3)
tensorrt_llm/_torch/compilation/backend.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Root cause analysis looks right, but the retirement is at the wrong granularity. Backend._graph_pool_handle is class-level and shared by every Backend/PyTorchModelEngine in the process, while retire_graph_pool_handle() is called from a per-engine _release_cuda_graphs(). If two engines coexist (e.g. draft + target model for spec decoding, or an encoder engine), tearing down one nulls the handle that the other's still-live graphs were captured into; the next capture in the surviving engine mints a fresh pool and silently loses memory sharing, and any code still holding the old handle in _cuda_graph_mem_pool passes a retired pool to capture_begin — the exact assert this PR fixes.
Either refcount the handle (number of engines that took it) and retire on the last release, or move the handle to per-Backend-instance state.
Also: the description says self._cuda_graph_mem_pool is nulled, but the diff never does that. And there's no regression test for the build → shutdown → rebuild sequence, which is the whole failure mode; the unwaived integration test only covers a single-engine run.
| # (e.g. a second LLM built after this one shuts down) passes it to | ||
| # capture_begin. | ||
| Backend.retire_graph_pool_handle() | ||
|
|
There was a problem hiding this comment.
This retires a class-level handle from a per-engine teardown. With more than one engine alive in the process (spec-decode draft+target, encoder+decoder), releasing one engine's graphs nulls the pool that the other engine's still-captured graphs own. The survivor's next capture then mints a new pool (losing sharing), and anything that cached the old value — e.g. self._cuda_graph_mem_pool on the other engine — now holds a handle whose graphs may have been reset.
Suggest refcounting: Backend.acquire_graph_pool_handle() on engine init, release_graph_pool_handle() here, and only null when the count hits zero.
| self.spec_metadata = None | ||
| self.iter_states = {} | ||
| self._cuda_graph_mem_pool = self._torch_compile_backend._graph_pool_handle if self._torch_compile_enabled else None | ||
| self._cuda_graph_mem_pool = (Backend.get_graph_pool_handle() |
There was a problem hiding this comment.
The PR description says _release_cuda_graphs() also nulls self._cuda_graph_mem_pool, but it doesn't. After retirement this field still holds the retired handle, and any subsequent capture path that reads it (rather than re-calling get_graph_pool_handle()) will pass the dead pool to capture_begin. Please either set it to None in _release_cuda_graphs() or make it a property that always delegates to the accessor.
| torch.cuda.Event() for _ in range(num_events - len(self.events)) | ||
| ] | ||
|
|
||
| @classmethod |
There was a problem hiding this comment.
No locking here. If two Backend instances are constructed or optimize concurrently (torch.compile worker threads, or the num_streams path), this can mint two pools and leak one. Probably fine in practice given the GIL and current call sites, but worth confirming that optimize()/engine init are single-threaded.
| @@ -246,7 +246,6 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4 | |||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-dp4-trtllm-fp8] SKIP (https://nvbugs/6474894) | |||
There was a problem hiding this comment.
Unwaiving is right if the fix holds, but the failure mode is a second LLM built in a reused MPI worker — a single-engine accuracy test won't catch a regression here. Worth adding a cheap unit test that constructs an engine with torch_compile, calls _release_cuda_graphs(), then asserts Backend.get_graph_pool_handle() returns a different handle and a fresh capture succeeds.
| return cls._graph_pool_handle | ||
|
|
||
| @classmethod | ||
| def retire_graph_pool_handle(cls) -> None: |
There was a problem hiding this comment.
clear_piecewise_cuda_graphs() already implements this same "the handle is dead, mint a new one" rule by calling torch.cuda.graph_pool_handle() itself, and the two mint different ids — so after a release the surviving runners sit on one pool while the next optimize() and the CUDAGraphRunner land on another, silently losing the backend-wide pool sharing this handle exists for.
Could clear_piecewise_cuda_graphs() instead call retire_graph_pool_handle() and then assign get_graph_pool_handle() to all runners? That keeps the rotation rule in one place and preserves the sharing.
9a4ea0a to
55ddcea
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
55ddcea to
d6fe9cf
Compare
…eleased Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
d6fe9cf to
1a2a18f
Compare
4 participants
Summary
Backend._graph_pool_handleis a process-wide class attribute never retired when_release_cuda_graphs()resets the graphs captured into it, so a later engine in the same reused MPI worker passes the now-dead private pool tocapture_beginand trips the allocator'suse_count > 0assert.Backend.get_graph_pool_handle()plusBackend.retire_graph_pool_handle()invoked at the end of_release_cuda_graphs()(also nullingself._cuda_graph_mem_pool), so the next capture generation mints a fresh handle; removed the obsolete waiver.pytest "tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=FLASHINFER-torch_compile=True]" "tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=True]" -vTest plan
Links
Dev Engineer Review
TestLlama3_1_8BInstruct::test_fp8withfp8kv=False, TRTLLM attention, and Torch Compile.QA Engineer Review
tests/integration/test_lists/waives.txt.